#include "SampleGraph.h"
#include "SharedFileName.h"
#include <fstream>
#include <sstream>
#include <iostream>
#include <stdlib.h>
#include <time.h>
#include <unordered_map>
#include <unordered_set>
#include <algorithm>
using namespace std;

// related with the seed from graph sample, if the graph node has a degree larger than this threshold, it can be used as a seed for breath first sampling
#define GRAPH_SAMPLE_NODE_DEGREE_THRESHOLD 2 

SampleGraph::SampleGraph(int sample_size, string in_node_file, string in_edge_file, string out_dir){
	//RebuildSampleGraph(out_dir); return;

	map<int,vector<int> > full_adj_list;
	vector<int> all_nodes;
	map<int, string> node_names;
	SampleNode(in_node_file,all_nodes,node_names);

	string out_node_file(out_dir);
	string out_adj_file(out_dir);
	out_node_file.append(sample_node_file);
	out_adj_file.append(sample_adjlist_file);

	SampleEdge(sample_size,in_edge_file,all_nodes,node_names,out_node_file,out_adj_file);
}

void SampleGraph::SampleNode(string in_node_file, vector<int>& all_nodes,  map<int,string>& node_names){
	ifstream myfile (in_node_file.c_str());
	int num_users = 0;

	if(myfile.is_open()){
		string line;
		while(getline(myfile,line)){
			stringstream ss(line);
			string user;
			string user_name;
			ss >> user;
			if(user.empty()) continue;
			ss >> user_name;
			if(user_name.empty()) continue;
			int user_id = atoi(user.c_str());
			
			// add the nodes to all possible nodes
			all_nodes.push_back(user_id);
			node_names[user_id] = user_name;
			num_users++;
			if(num_users % 1000000 == 0) cout << "read " << (num_users / 1000000) << "M users" << endl;
		}
		myfile.close();
	}else{
		cout << "cannot read node file." << endl;
		exit(1);
	}

	cout << "finish reading node file with " << num_users << " users" << endl;
}

void SampleGraph::SampleEdge(int sample_size, string in_edge_file, vector<int>& all_nodes, map<int,string>& node_names,string out_node_file, string out_adj_file){

	unordered_map<int,int> node_idx_map; // index map for the nodes
	for(int i = 0; i < all_nodes.size(); i++){
		node_idx_map[all_nodes[i]] = i;
	}
	vector<vector<int> > full_adj_list_idx(all_nodes.size()); // full adjacency list index, for easy find neighbours

	// extract the edges
	long long total_edge = 0;
	ifstream edge_file(in_edge_file.c_str());
	int max_length = 0;
	int max_following_id;
	if(edge_file.is_open()){
		while(true){
			int followed_id,following_id;
			if(!(edge_file >> followed_id)) break;
			if(!(edge_file >> following_id)) break;

			if(followed_id == following_id) continue;

			unordered_map<int,int>::iterator it = node_idx_map.find(followed_id);
			if(it == node_idx_map.end()) continue;
			it = node_idx_map.find(following_id);
			if(it == node_idx_map.end()) continue;
			int following_idx = it->second;

			full_adj_list_idx[following_idx].push_back(followed_id);
			if(full_adj_list_idx[following_idx].size() > max_length){
				max_length = full_adj_list_idx[following_idx].size();
				max_following_id = following_id;
			}

			total_edge++;
			if(total_edge % 10000000 == 0) cout << "read " << (total_edge / 10000000) << "x10M edges with max length: " << max_length << " and the ID is: " << max_following_id << endl;
		}
		edge_file.close();
	}else{
		cout << "cannot read edge file." << endl;
		exit(1);
	}
	cout << "finish reading edge file" << endl;


	for(int i = 0; i < full_adj_list_idx.size(); i++){
		sort(full_adj_list_idx[i].begin(),full_adj_list_idx[i].end());
		auto it = unique(full_adj_list_idx[i].begin(),full_adj_list_idx[i].end());
		full_adj_list_idx[i].resize(std::distance(full_adj_list_idx[i].begin(),it));
	}
	cout << "finish sorting vector" << endl;

	vector<vector<int > > full_adj_list(all_nodes.size());
	
	// filter out the friends relationship (follow each other)
	int num_undirected_edges = 0;
	for(int i = 0; i < full_adj_list_idx.size(); i++){
		for(int j = 0; j < full_adj_list_idx[i].size(); j++){
			int neighbour = full_adj_list_idx[i][j];
			if(neighbour < all_nodes[i]) continue; // if the neighbor index is small than the current node, this means the edge has been computed
			int neighbor_idx = node_idx_map[full_adj_list_idx[i][j]];
			if(binary_search(full_adj_list_idx[neighbor_idx].begin(),full_adj_list_idx[neighbor_idx].end(),all_nodes[i])){
				full_adj_list[i].push_back(neighbour);
				full_adj_list[neighbor_idx].push_back(all_nodes[i]);
				num_undirected_edges++;
			}	
		}
		full_adj_list_idx[i].clear();
		if(i % 1000000 == 0) cout << "extracted " << (i / 1000000) << "M nodes that has bi-directional relationship" << endl;
	}
	full_adj_list_idx.clear(); // clear memory
	cout << "finish extracting bi-directional relationship and the graph has " << num_undirected_edges << " edges"<< endl;
	
	double degree_sum = 0;
	for(int i = 0; i < all_nodes.size(); i++) degree_sum += full_adj_list[i].size();
	cout << "total degree: " << degree_sum << endl;
	cout << "average degree for full friendship (undirected graph is) " << degree_sum / all_nodes.size() << endl;

	// use BFS to random sample connected nodes which has size sample_size 
	int number_sampled = 0;	// keep track of the current connected component size
	unordered_set<int> sample_froniter;	// keep track of sample froniter
	vector<bool> visited_nodes(all_nodes.size(),false); // for all BFS sampling, keep all the visited nodes
	unordered_set<int> connected_nodes;	// this is the sampled results

	while(number_sampled < sample_size){
		unordered_set<int> new_froniter; // update frontier

		if(sample_froniter.empty()){
			cout << "sampling a new connected component" << endl;
			connected_nodes.clear(); // restart the sampling

			bool all_visited = true;
			for(int i = 0; i < visited_nodes.size(); i++){
				if(!visited_nodes[i] && full_adj_list[i].size() > GRAPH_SAMPLE_NODE_DEGREE_THRESHOLD){
					all_visited = false;
					number_sampled = 1;

					visited_nodes[i] = true; // mark this node has been sampled before
					connected_nodes.insert(all_nodes[i]);
					new_froniter.insert(all_nodes[i]);

					break;
				}
			}
			if(all_visited){
				cout << "problem with the BFS sampling" << endl;
				break;
			}
		}else{
			for(auto it = sample_froniter.begin(); it != sample_froniter.end(); it++){
				vector<int>& neighbour_vec = full_adj_list[node_idx_map[*it]];
				for(int i = 0; i < neighbour_vec.size(); i++){
					int neighbour_idx = node_idx_map[neighbour_vec[i]];  // index start from 0 for the full adjacency list
					if(!visited_nodes[neighbour_idx]){
						number_sampled++;

						visited_nodes[neighbour_idx] = true;
						connected_nodes.insert(neighbour_vec[i]);
						new_froniter.insert(neighbour_vec[i]);

						if(number_sampled == sample_size) break;
					}
				}
				if(number_sampled == sample_size) break;
			}
		}
		cout << "sampled " << number_sampled << " nodes" << endl;
		sample_froniter = new_froniter;
	}
	cout << "finish sampling nodes" << endl;


	map<int,string> sampled_names; // store the names that has been sampled
	for(auto it = connected_nodes.begin(); it != connected_nodes.end(); it++) sampled_names[*it] = node_names[*it];
	
	vector<bool> sampled_indicator(all_nodes.size(),false);
	unordered_map<int,int> sampled_node_map; // convert the index from 0 and continuously

	// write the sampled nodes to file
	ofstream nodeFile(out_node_file.c_str(),ios::out);
	int inc_idx = 0;
	for(auto it = sampled_names.begin(); it != sampled_names.end(); it++){
		nodeFile << it->second << endl;
		sampled_indicator[node_idx_map[it->first]] = true;
		sampled_node_map[it->first] = inc_idx;
		inc_idx++;
	}
	nodeFile.close();
	cout << "finish write sampled nodes with size " << inc_idx << endl;

	// get the edge in the sampled set
	vector<vector<int> > sampled_adj_list(sample_size);
	for(auto it = sampled_names.begin(); it != sampled_names.end(); it++){
		int node_ID = it->first;

		vector<int>& neighbours = full_adj_list[node_idx_map[node_ID]];
		for(int i = 0; i < neighbours.size(); i++){
			if(neighbours[i] < node_ID) continue;
			if(sampled_indicator[node_idx_map[neighbours[i]]]){
				sampled_adj_list[sampled_node_map[node_ID]].push_back(sampled_node_map[neighbours[i]]);
				sampled_adj_list[sampled_node_map[neighbours[i]]].push_back(sampled_node_map[node_ID]);
			}
		}
	}
	cout << "finish finding the edges" << endl;

	int number_edges = 0;
	ofstream myfile(out_adj_file.c_str(),ios::out); // without appending to avoid mistake
	
	// write number of nodes
	myfile << sample_size << endl;

	// each line write the following nodes
	for(auto it = sampled_names.begin(); it != sampled_names.end(); it++){
		int mapped_sample_idx = sampled_node_map[it->first];

		// write number of adjacent nodes
		int list_size = sampled_adj_list[mapped_sample_idx].size();
		myfile << list_size;
		number_edges += list_size;

		// write the adjacent nodes
		for(int i = 0; i < list_size; i++)
			myfile << " " << sampled_adj_list[mapped_sample_idx][i];
		
		myfile << endl;
	}

	myfile.close();
	cout << "sampled " << sample_size << " nodes with " << number_edges/2 << " edges" << endl;
}

int main(int argc, char *argv[]){
	if(argc != 5){
		cout << "wrong arguments for " << argv[0] << endl;
		cout << "1. number of nodes that you want to sample" << endl;
		cout << "2. the node information (userID + userName)" << endl;
		cout << "3. relationship graph" << endl;
		cout << "4. output directory" << endl;
		exit(1);
	}

	int sample_size = atoi(argv[1]);
	string in_node_file(argv[2]);
	string in_edge_file(argv[3]);
	string out_dir(argv[4]);

	cout << "args: " << sample_size << " " << in_node_file << " " << in_edge_file << " " << out_dir << endl;

	SampleGraph sample_graph(sample_size,in_node_file,in_edge_file,out_dir);
	cout << argv[0] << " finished" << endl;
}